package org.fujay.commons.langchain4j.rag.prompt.builder;

import dev.langchain4j.rag.content.Content;
import org.fujay.commons.langchain4j.rag.prompt.template.ChainOfThoughtTemplate;
import org.fujay.commons.langchain4j.rag.prompt.template.ExpertRoleTemplate;
import org.fujay.commons.langchain4j.rag.prompt.template.RagPromptTemplate;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 默认提示构建器
 * 根据不同的场景选择不同的提示模板
 */
public class DefaultPromptBuilder implements PromptBuilder {

    private static final String DEFAULT_TEMPLATE = """
            请基于以下信息回答问题。如果检索到的信息不足以回答问题，请诚实地表示不知道。
            
            检索到的信息：
            %s
            
            用户问题: %s""";

    private final RagPromptTemplate defaultTemplate;
    private final Map<String, RagPromptTemplate> templateMap;

    /**
     * 构造函数
     */
    public DefaultPromptBuilder() {
        this(null, null);
    }

    /**
     * 构造函数
     *
     * @param defaultTemplate 默认模板
     * @param templateMap     模板映射
     */
    public DefaultPromptBuilder(RagPromptTemplate defaultTemplate, Map<String, RagPromptTemplate> templateMap) {
        this.defaultTemplate = defaultTemplate;
        this.templateMap = templateMap != null ? new HashMap<>(templateMap) : new HashMap<>();

        // 如果没有提供默认模板，使用简单模板
        if (this.defaultTemplate == null) {
            final String template = DEFAULT_TEMPLATE;
            this.templateMap.put("default", (query, contents) -> {
                StringBuilder formattedContents = new StringBuilder();
                for (int i = 0; i < contents.size(); i++) {
                    formattedContents.append(i + 1).append(". ").append(contents.get(i).toString()).append("\n\n");
                }
                return String.format(template, formattedContents.toString(), query);
            });
        }

        // 添加内置模板
        if (!this.templateMap.containsKey("chain-of-thought")) {
            this.templateMap.put("chain-of-thought", new ChainOfThoughtTemplate());
        }
    }

    /**
     * 使用指定模板构建提示词
     *
     * @param templateKey       模板键
     * @param query             查询
     * @param retrievedContents 检索到的内容
     * @return 构建好的提示词
     */
    public String build(String templateKey, String query, List<Content> retrievedContents) {
        // 获取指定模板，如果不存在则使用默认模板
        RagPromptTemplate template = templateMap.getOrDefault(templateKey,
                templateMap.getOrDefault("default", defaultTemplate));

        if (template == null) {
            // 如果仍然没有模板，使用最简单的格式拼接
            StringBuilder result = new StringBuilder("用户问题: ").append(query).append("\n\n检索到的信息:\n");
            for (int i = 0; i < retrievedContents.size(); i++) {
                result.append(i + 1).append(". ").append(retrievedContents.get(i).toString()).append("\n");
            }
            return result.toString();
        }

        return template.build(query, retrievedContents);
    }

    @Override
    public String build(String query, List<Content> retrievedContents) {
        return build("default", query, retrievedContents);
    }

    /**
     * 使用思维链模板构建提示词
     *
     * @param query             查询
     * @param retrievedContents 检索到的内容
     * @return 构建好的提示词
     */
    public String buildWithChainOfThought(String query, List<Content> retrievedContents) {
        return build("chain-of-thought", query, retrievedContents);
    }

    /**
     * 使用专家角色模板构建提示词
     *
     * @param domain            专业领域
     * @param query             查询
     * @param retrievedContents 检索到的内容
     * @return 构建好的提示词
     */
    public String buildWithExpertRole(String domain, String query, List<Content> retrievedContents) {
        // 动态创建专家角色模板
        ExpertRoleTemplate expertTemplate = new ExpertRoleTemplate(domain);
        return expertTemplate.build(query, retrievedContents);
    }

    /**
     * 添加模板
     *
     * @param key      模板键
     * @param template 模板
     * @return this
     */
    public DefaultPromptBuilder addTemplate(String key, RagPromptTemplate template) {
        if (key != null && template != null) {
            templateMap.put(key, template);
        }
        return this;
    }
}